package com.ai.module.chat.util;

import cn.hutool.extra.spring.SpringUtil;
import com.ai.module.chat.common.base.ThreadLocalContext;
import com.ai.module.chat.common.enu.ModelType;
import com.ai.module.chat.config.properties.ChatGptProperty;
import com.ai.module.websocket.server.WebSocketController;
import com.ai.module.business.domain.vo.UserMessage;
import com.ai.module.websocket.service.AIProductWebSocketService;
import com.ai.module.websocket.service.impl.AIProductWebSocketServiceImpl;
import com.ai.module.websocket.vo.WebSocketProduct;
import com.alibaba.fastjson2.JSONObject;
import com.ai.module.user.domain.req.DallECreateImageRequest;
import com.ai.module.user.domain.vo.UserTimeVo;
import com.baomidou.mybatisplus.core.toolkit.StringPool;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.theokanning.openai.completion.chat.ChatCompletionChoice;
import com.theokanning.openai.completion.chat.ChatCompletionChunk;
import com.theokanning.openai.completion.chat.ChatCompletionRequest;
import com.theokanning.openai.completion.chat.ChatMessage;
import com.theokanning.openai.image.CreateImageRequest;
import io.github.asleepyfish.config.ChatGPTProperties;
import io.github.asleepyfish.enums.chat.RoleEnum;
import io.github.asleepyfish.enums.exception.ChatGPTErrorEnum;
import io.github.asleepyfish.enums.image.ImageResponseFormatEnum;
import io.github.asleepyfish.enums.image.ImageSizeEnum;
import io.github.asleepyfish.exception.ChatGPTException;
import io.github.asleepyfish.service.OpenAiProxyService;
import lombok.extern.slf4j.Slf4j;
import okhttp3.OkHttpClient;
import org.jetbrains.annotations.NotNull;
import org.springframework.beans.BeanUtils;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.socket.WebSocketSession;

import javax.imageio.ImageIO;
import javax.websocket.Session;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import static com.ai.module.chat.common.constant.StringConstant.IMAGE_URL;
import static com.ai.module.chat.common.constant.StringConstant.userROBOT;


/**
 * @author: MrHe
 * @create: 2023-11-14 21:38
 * @description:
 **/
@Slf4j
public class MrHeOpenAiService extends OpenAiProxyService {

    private final ChatGPTProperties chatGPTProperties;

    private final OkHttpClient client;

    private Cache<String, LinkedList<ChatMessage>> cache;

    private static final Random RANDOM = new Random();

    public MrHeOpenAiService(ChatGPTProperties chatGPTProperties, Duration timeout) {
        super(chatGPTProperties);
        this.chatGPTProperties = chatGPTProperties;
        this.cache = chatGPTProperties.getSessionExpirationTime() == null ? CacheBuilder.newBuilder().build() :
                CacheBuilder.newBuilder().expireAfterAccess(chatGPTProperties.getSessionExpirationTime(), TimeUnit.MINUTES).build();
        this.client = OpenAiProxyService.defaultClient(chatGPTProperties, timeout);
    }


    public void createStreamChat4(UserTimeVo userTimeVo, OutputStream os) {
        userTimeVo.setRole(RoleEnum.USER.getRoleName());
        ChatGptProperty chatGptProperty = getChatModel4();
        createStreamChatCompletion( userTimeVo,  chatGptProperty.getChatModel4(), 1.0D, 1.0D, os);
    }


    public void createStreamChatCompletion(UserTimeVo userTimeVo,
                                           String model,
                                           Double temperature,
                                           Double topP,
                                           OutputStream os) {
        createStreamChatCompletion(ChatCompletionRequest.builder()
                .model(model)
                .messages(Collections.singletonList(new ChatMessage(userTimeVo.getRole(), userTimeVo.getAskContent())))
                .user(userTimeVo.getUser())
                // 较高的值(如 0.8)会使输出更随机，而较低的值(如 0.2)则会使其更加专注和确定。
                .temperature(temperature)
                .topP(topP)
                .stream(true)
                .build(),userTimeVo, os);
    }

    public void createStreamChatCompletion(ChatCompletionRequest chatCompletionRequest,
                                           UserTimeVo userTimeVo,
                                           OutputStream os) {
        String askMessage = chatCompletionRequest.getMessages().get(0).getContent();
        StringBuilder sb = new StringBuilder();
        chatCompletionRequest.setStream(true);
        chatCompletionRequest.setN(1);
        String user = chatCompletionRequest.getUser();
        LinkedList<ChatMessage> contextInfo = new LinkedList<>();
        try {
            contextInfo = cache.get(user, LinkedList::new);
        } catch (ExecutionException e) {
            e.printStackTrace();
        }
        contextInfo.addAll(chatCompletionRequest.getMessages());
        chatCompletionRequest.setMessages(contextInfo);
        List<ChatCompletionChunk> chunks = new ArrayList<>();
        for (int i = 0; i < chatGPTProperties.getRetries(); i++) {
            try {
                // avoid frequently request, random sleep 0.5s~0.7s
                if (i > 0) {
                    randomSleep();
                }
                super.streamChatCompletion(chatCompletionRequest).doOnError(Throwable::printStackTrace).blockingForEach(chunk -> {
                    chunk.getChoices().stream().map(choice -> choice.getMessage().getContent())
                            .filter(Objects::nonNull).findFirst().ifPresent(o -> {
                                try {
                                    if (userTimeVo.getBeginTime() == 0){
                                        userTimeVo.setBeginTime(System.currentTimeMillis());
                                    }
                                    ObjectMapper objectMapper = new ObjectMapper();

                                    String jsonResponse = objectMapper.writeValueAsString(o);
                                   // jsonResponse = jsonResponse
                                   //         .replace(BaseSymbolConstant.LINE_JAVA, BaseSymbolConstant.NEW_LINE)
                                   //         .replace(BaseSymbolConstant.LINE_JAVA, BaseSymbolConstant.NEW_LINE)
                                   //         .replace(BaseSymbolConstant.BLANK_JAVA, BaseSymbolConstant.BLANK)
                                   //         .replaceAll(BaseSymbolConstant.BLANK_JAVA, BaseSymbolConstant.BLANK)
                                   //         .replaceAll(BaseSymbolConstant.SYMBOL_JAVA, BaseSymbolConstant.SYMBOL)
                                   // ;
                                    String sseFormattedResponse = "data: " + jsonResponse + "\n\n";
                                    os.write(sseFormattedResponse.getBytes(StandardCharsets.UTF_8));
                                    os.flush();
                                    sb.append(jsonResponse);
//                                    os.write(o.getBytes(Charset.defaultCharset()));
//                                    os.flush();
                                } catch (Exception e) {
                                    e.printStackTrace();
                                    throw new RuntimeException(e);
                                }
                            });
                    chunks.add(chunk);
                });
                os.write(( "data: " +"end"+ "\n\n").getBytes(StandardCharsets.UTF_8));
                os.flush();
                os.close();
                // if the last line code is correct, we can simply break the circle
                break;
            } catch (Exception e) {
                String message = e.getMessage();
                boolean overload = checkTokenUsage(message);
                if (overload) {
                    int size = Objects.requireNonNull(cache.getIfPresent(user)).size();
                    for (int j = 0; j < size / 2; j++) {
                        ChatMessage removeMessage = Objects.requireNonNull(cache.getIfPresent(user)).removeFirst();
                    }
                    chatCompletionRequest.setMessages(cache.getIfPresent(user));
                }
                log.error("answer failed " + (i + 1) + " times, the error message is: " + message);
                if (i == chatGPTProperties.getRetries() - 1) {
                    e.printStackTrace();
                    // when the call fails, remove the last item in the list
                    Objects.requireNonNull(cache.getIfPresent(user)).removeLast();
                    // 进行数据处理?
                    throw new ChatGPTException(ChatGPTErrorEnum.FAILED_TO_GENERATE_ANSWER, message);
                }

            }
        }
        LinkedList<ChatMessage> chatMessages = new LinkedList<>();
        try {
            chatMessages = cache.get(user, LinkedList::new);
        } catch (ExecutionException e) {
            e.printStackTrace();
        }
        chatMessages.add(new ChatMessage(RoleEnum.ASSISTANT.getRoleName(), chunks.stream()
                .flatMap(chunk -> chunk.getChoices().stream())
                .map(ChatCompletionChoice::getMessage)
                .map(ChatMessage::getContent)
                .filter(Objects::nonNull)
                .collect(Collectors.joining())));
        userTimeVo.setChatGptTalk(sb.toString());
        userTimeVo.setLastTime(System.currentTimeMillis());
        userTimeVo.setAskContent(askMessage);
    }

    private void randomSleep() throws InterruptedException {
        Thread.sleep(1000 + RANDOM.nextInt(100));
    }
    public void createStreamChat4WebSocket(UserTimeVo userTimeVo,
                                           WebSocketController webSocketController,
                                           Session session) {
        boolean test =false;
        if (false){
            userTimeVo.setChatGptTalk("loadCookie-");
            return;
        }
        userTimeVo.setRole(RoleEnum.USER.getRoleName());
        ChatGptProperty chatGptProperty = getChatModel4();
        String model = chatGptProperty.getChatModel3();
        ModelType modelType = userTimeVo.getModelType();
        if (modelType == ModelType.CHAT_4){
            model= chatGptProperty.getChatModel4();
        }
        createStreamChatCompletionWebSocket( userTimeVo,  model,
                1.0D, 1.0D, webSocketController,session);
    }
    public void createStreamChatAIProductWebSocket( UserTimeVo userTimeVo ) {
        boolean test =false;
        if (false){
            userTimeVo.setChatGptTalk("loadCookie-");
            return;
        }
        userTimeVo.setRole(RoleEnum.USER.getRoleName());
        ChatGptProperty chatGptProperty = getChatModel4();
        String model = chatGptProperty.getChatModel3();
        ModelType modelType = userTimeVo.getModelType();
        if (modelType == ModelType.CHAT_4){
            model= chatGptProperty.getChatModel4();
        }
        createStreamChatCompletionWebSocket( userTimeVo,  model,
                1.0D, 1.0D, null,null);
    }

    private void createStreamChatCompletionWebSocket(UserTimeVo userTimeVo,
                                                     String model,
                                                     double temperature,
                                                     double topP,
                                                     WebSocketController webSocketController,
                                                     Session session) {
        createStreamChatCompletionWebSocket(ChatCompletionRequest.builder()
                .model(model)
                .messages(Collections.singletonList(new ChatMessage(userTimeVo.getRole(), userTimeVo.getAskContent())))
                .user(userTimeVo.getUser() + StringPool.AT+DateUtils.getPaYear())
                .temperature(temperature)
                .topP(topP)
                .stream(true)
                .build(),userTimeVo, webSocketController,session);
    }



    public void createStreamChatCompletionWebSocket(ChatCompletionRequest chatCompletionRequest,
                                                    UserTimeVo userTimeVo,
                                                    WebSocketController webSocketController,
                                                    Session session) {
        String askMessage = chatCompletionRequest.getMessages().get(0).getContent();
        //StringBuilder sb = new StringBuilder();
        AIProductWebSocketService globalWebSocketService = SpringUtil.getBean(AIProductWebSocketService.class);
        Map<String, WebSocketSession> userProductSessionMap = AIProductWebSocketServiceImpl.userProductSessionMap;
        String userUUID = userTimeVo.getUser();
        String robotId = userTimeVo.getRobotId();
        String userIdKey = String.format(userROBOT, userUUID, robotId);
        WebSocketSession webSocketSession = userProductSessionMap.get(userIdKey);
//        BA792941-15C5-45E2-B207-2800521A6B2E@5D839660-2016-4D77-8214-425A5DDC5079
//        BA792941-15C5-45E2-B207-2800521A6B2E@5d839660-2016-4d77-8214-425a5ddc5079 -> {StandardWebSocketSession@13761} "StandardWebSocketSession[id=bca5a88a-e1dd-3673-1754-86348e0ffa4f, uri=ws://127.0.0.1:18080/websocket/aiProductChat]"
        chatCompletionRequest.setStream(true);
        chatCompletionRequest.setN(1);
        String user = chatCompletionRequest.getUser();
        StringBuilder sb = new StringBuilder();
        UserMessage userMessage = new UserMessage();
        WebSocketProduct webSocketProduct = new WebSocketProduct();
        WebSocketProduct.AiChat aiChat = new WebSocketProduct.AiChat();
        webSocketProduct.setAiChat(aiChat);
        BeanUtils.copyProperties(userTimeVo, userMessage);
        LinkedList<ChatMessage> contextInfo = new LinkedList<>();
        boolean isBanner= false;
        try {
            contextInfo = cache.get(user, LinkedList::new);
        } catch (ExecutionException e) {
            e.printStackTrace();
        }
        contextInfo.addAll(chatCompletionRequest.getMessages());
        chatCompletionRequest.setMessages(contextInfo);
        List<ChatCompletionChunk> chunks = new ArrayList<>();
        for (int i = 0; i < chatGPTProperties.getRetries(); i++) {
            try {
                // avoid frequently request, random sleep 0.5s~0.7s
                if (i > 0) {
                    //重置用户
                    randomSleep();
                }
                super.streamChatCompletion(chatCompletionRequest).doOnError(Throwable::printStackTrace).blockingForEach(chunk -> {
                    chunk.getChoices().stream().map(choice -> choice.getMessage().getContent())
                            .filter(Objects::nonNull).findFirst().ifPresent(o -> {
                                try {
                                    if (userTimeVo.getBeginTime() == 0){
                                        userTimeVo.setBeginTime(System.currentTimeMillis());
                                    }
                                    ObjectMapper objectMapper = new ObjectMapper();

                                    String jsonResponse = objectMapper.writeValueAsString(o);
                                    String originalRe = jsonResponse;
                                    jsonResponse = jsonResponse
                                            .replaceAll("\"\"", "")
                                            .replace("\\\"", "\"")
                                            .replaceAll("\"", "");

                                    //   .replaceAll(BaseSymbolConstant.SYMBOL_COMPLEX, BaseSymbolConstant.BLANK)
                                  //   .replaceAll(BaseSymbolConstant.SYMBOL_JAVA_F_2, BaseSymbolConstant.BLANK)
                                    ;

                                    sb.append(originalRe);
                                    String regex = "(?i)(chatgpt|openai|gpt)";
                                    String answerContent = sb.toString();
                                    // 检查字符串是否匹配
                                    boolean containsSensitiveWords = answerContent.matches(".*" + regex + ".*");
//                                    判断是不是 ai-product

                                    if (containsSensitiveWords){
                                        String message ="***";
                                        if (webSocketSession ==null){
                                            userMessage.setUserMessage(message);
                                            webSocketController.sendMessage(session, userMessage);
                                        } else {
                                            aiChat.setMessage(message);
                                            globalWebSocketService.sendMessage(webSocketSession,webSocketProduct);
                                        }
                                        // sb.append(jsonResponse);
                                 }else {
                                        if (webSocketSession ==null){
                                            userMessage.setUserMessage(jsonResponse);
                                            webSocketController.sendMessage(session, userMessage);
                                        } else {
                                            aiChat.setMessage(jsonResponse);
                                            globalWebSocketService.sendMessage(webSocketSession,webSocketProduct);
                                        }

                                    }
                                    // sb.append(jsonResponse);

//                                    os.write(o.getBytes(Charset.defaultCharset()));
//                                    os.flush();
                                } catch (Exception e) {
                                    throw new RuntimeException(e);
                                }
                            });
                    chunks.add(chunk);
                });
                if (webSocketSession ==null){
                    userMessage.setUserMessage(StringPool.EMPTY);
                    userMessage.setEnd(true);
                    webSocketController.sendMessage(session, userMessage);
                } else {
                    aiChat.setMessage(StringPool.EMPTY);
                    aiChat.setEnd(true);
                    globalWebSocketService.sendMessage(webSocketSession,webSocketProduct);
                }

                // if the last line code is correct, we can simply break the circle
                break;
            } catch (Exception e) {
                String message = e.getMessage();
                boolean overload = checkTokenUsage(message);
                if (overload) {
                    int size = Objects.requireNonNull(cache.getIfPresent(user)).size();
                    for (int j = 0; j < size / 2; j++) {
                        ChatMessage removeMessage = Objects.requireNonNull(cache.getIfPresent(user)).removeFirst();
                    }
                    chatCompletionRequest.setMessages(cache.getIfPresent(user));
                }
                log.error("answer failed " + (i + 1) + " times, the error message is: " + message);
                if (i == chatGPTProperties.getRetries() - 1) {
                    e.printStackTrace();
                    // when the call fails, remove the last item in the list
                    Objects.requireNonNull(cache.getIfPresent(user)).removeLast();
                    // 进行数据处理?
                    throw new ChatGPTException(ChatGPTErrorEnum.FAILED_TO_GENERATE_ANSWER, message);
                }

            }
        }
        LinkedList<ChatMessage> chatMessages = new LinkedList<>();
        try {
            chatMessages = cache.get(user, LinkedList::new);
        } catch (ExecutionException e) {
            e.printStackTrace();
        }
        chatMessages.add(new ChatMessage(RoleEnum.ASSISTANT.getRoleName(), chunks.stream()
                .flatMap(chunk -> chunk.getChoices().stream())
                .map(ChatCompletionChoice::getMessage)
                .map(ChatMessage::getContent)
                .filter(Objects::nonNull)
                .collect(Collectors.joining())));
        userTimeVo.setChatGptTalk( sb.toString());
        userTimeVo.setLastTime(System.currentTimeMillis());
        userTimeVo.setAskContent(askMessage);
    }

    private static boolean checkTokenUsage(String message) {
        return message != null && message.contains("This model's maximum context length is");
    }


    public static BufferedImage getImageFromBase64(String base64) throws IOException {
        byte[] imageBytes = Base64.getDecoder().decode(base64.getBytes());
        try (ByteArrayInputStream bis = new ByteArrayInputStream(imageBytes)) {
            return ImageIO.read(bis);
        }
    }

    /**
     *  生成imageR对象
     * @param desc
     * @return 图片描述
     */
    public static CreateImageRequest createImageRequest(String desc) {
        CreateImageRequest request = CreateImageRequest.builder()
                .prompt(desc)
                .n(1)
                .size(ImageSizeEnum.S1024x1024.getSize())
                .user("DEFAULT USER").build();
        DallECreateImageRequest dallECreateImageRequest = new DallECreateImageRequest();
        BeanUtils.copyProperties(request,dallECreateImageRequest);
        dallECreateImageRequest.setModel(null);
        dallECreateImageRequest.setResponseFormat(ImageResponseFormatEnum.B64_JSON.getResponseFormat());
        return dallECreateImageRequest;
    }

    public static JSONObject createImage4(String desc) {
        long start = System.currentTimeMillis();
        ChatGptProperty chat4Model = getChatModel4();
        RestTemplate restTemplate = SpringUtils.getClazz(RestTemplate.class,"restTemplate");
        String dall = chat4Model.getDall();
        String baseUrl = chat4Model.getBaseUrl() + IMAGE_URL;
        CreateImageRequest default_user = CreateImageRequest.builder()
                .prompt(desc)
                .n(1)
                .size(ImageSizeEnum.S1024x1024.getSize() )
                .user(ThreadLocalContext.getUserId())
                .build();
        DallECreateImageRequest imageRequest4 = new DallECreateImageRequest();
        BeanUtils.copyProperties(default_user,imageRequest4);
        imageRequest4.setModel(dall);
        String requestJson = JSONObject.toJSONString(imageRequest4);
        HttpHeaders headers = new HttpHeaders();
        //headers.setContentType(MediaType.APPLICATION_JSON);
//        headers.set(HttpHeaders.AUTHORIZATION, "Bearer "+chat4Model.getVersion());
        headers.setBearerAuth(chat4Model.getToken());

        // 创建 HttpEntity
        HttpEntity<JSONObject> entity = new HttpEntity<>(JSONObject.parse(requestJson), headers);

        return restTemplate.postForEntity(baseUrl, entity, JSONObject.class).getBody();
       // assert jsonObject != null;
       // List<LinkedHashMap> mapList = jsonObject.getJSONArray("data").toJavaList(LinkedHashMap .class);
       // LinkedHashMap map= mapList.get(0);
       // return (String)map.get("url");
    }

    @NotNull
    private static ChatGptProperty getChatModel4() {
        ChatGptProperty chat4Model = SpringUtils.getClazz(ChatGptProperty.class);
        return chat4Model;
    }

    public List<String> chatCompletion(String content,
                                       String user,
                                       String model) {
        String role = RoleEnum.USER.getRoleName();
        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model(model)
                .messages(Collections.singletonList(new ChatMessage(role, content)))
                .user(user)
                .temperature(1.0D)
                .topP(1.0D)
                .build();
        return chatCompletion(request);
    }

    public List<String> chatCompletion(ChatCompletionRequest chatCompletionRequest) {
        return super.chatCompletion(chatCompletionRequest);
    }


}
